import cv2
from dataset import *
from model.FCN import Resnet101_FCN8s
from PIL import Image
import torchvision.transforms.functional as ff
from model.UNet import ResNet101_UNet


# device = 'cuda' if torch.cuda.is_available() else 'cpu'
test_image = './result/test'
transfrom = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]
    )
Cam_test_dataset = CamvidDataset([test_image,None],is_train=False,transform=transfrom)
test_dataloader = DataLoader(Cam_test_dataset,batch_size=1,shuffle=True,num_workers=0)
model = ResNet101_UNet(12)
model.load_state_dict(torch.load("./chekpoints/Unetbest.pth",map_location='cpu'))
model.eval()

label_color =pd.read_csv('./CamVid/class_dict.csv')
name_value = label_color['name'].values
num_class = len(name_value)
colormap = []
for i in range(num_class):
    temp = label_color.iloc[i]
    color = [temp['r'],temp['g'],temp['b']]
    colormap.append(color)
cm = np.array(colormap).astype('uint8')
dir = './result'
if not os.path.exists(dir):
    os.mkdir(dir)
for i,image in enumerate(test_dataloader):
    image_tensor = image
    out = model(image_tensor)
    print(out.shape)
    pre_label = out.max(1)[1].squeeze().cpu().data.numpy()
    pre = cm[pre_label]
    pre = Image.fromarray(pre)
    pre.save(dir+'/{}.png'.format(str(i)))
print('Done')





